import random
import numpy as np
import os
os.environ["MXNET_CPU_FORCE_F16"] = "1"

def fp16_to_hex(f):
    """将FP16数值转换为16进制字符串"""
    return f"{np.uint16(f):04X}"

def fp32_to_hex(f):
    """将FP32数值转换为32位十六进制字符串"""
    return f"{np.float32(f).view(np.uint32).item():08X}"

def generate_fp16_case():
    """生成单个FP16测试用例（包含NaN）"""
    num_type = random.choice(["normal", "denormal", "zero", "inf", "nan"])
    
    if num_type == "zero":
        sign = random.choice([0, 1])
        return (sign << 15) | 0x0000
    elif num_type == "inf":
        sign = random.choice([0, 1])
        return (sign << 15) | 0x7C00
    elif num_type == "denormal":
        sign = random.choice([0, 1])
        mantissa = random.randint(1, 0x3FF)
        return (sign << 15) | mantissa
    elif num_type == "nan":
        sign = random.choice([0, 1])
        mantissa = random.randint(1, 0x3FF)  # NaN尾数非零
        return (sign << 15) | 0x7C00 | mantissa
    else:  # normal
        sign = random.choice([0, 1])
        exp = random.randint(1, 0x1E)
        mantissa = random.randint(0, 0x3FF)
        return (sign << 15) | (exp << 10) | mantissa

def mixed_precision_mult(a, b):
    """混合精度乘法：显式处理Inf×0=NaN"""
    a_fp16 = np.uint16(a).view(np.float16)
    b_fp16 = np.uint16(b).view(np.float16)
    
    # 处理Inf × 0 = NaN
    if (np.isinf(a_fp16) and (b_fp16 == 0)) or ((a_fp16 == 0) and np.isinf(b_fp16)):
        return np.float32(np.nan)
    
    return a_fp16.astype(np.float32) * b_fp16.astype(np.float32)

# 生成测试用例（包含手动添加的关键用例）
test_cases = []
for _ in range(95):  # 随机生成95个用例
    a = generate_fp16_case()
    b = generate_fp16_case()
    result_fp32 = mixed_precision_mult(a, b)
    test_cases.append({
        "a": fp16_to_hex(a),
        "b": fp16_to_hex(b),
        "expected": fp32_to_hex(result_fp32)
    })

# 手动添加5个关键用例
key_cases = [
    {"a": "3B9A", "b": "8000", "expected": fp32_to_hex(-0.0)},  # 正数 × -0 = -0
    {"a": "7C00", "b": "0000", "expected": fp32_to_hex(np.nan)},  # Inf × 0 = NaN
    {"a": "0001", "b": "0001", "expected": fp32_to_hex(np.finfo(np.float16).tiny**2)},  # 非规格化数相乘
    {"a": "7C00", "b": "7C00", "expected": fp32_to_hex(np.inf)},  # Inf × Inf = Inf
    {"a": "8000", "b": "8000", "expected": fp32_to_hex(0.0)},     # -0 × -0 = +0
]
test_cases.extend(key_cases)

# 输出到文件
with open("fp16_to_fp32_mult_cases.txt", "w") as f:
    f.write("// FP16输入，FP32输出的混合精度乘法测试用例\n")
    for case in test_cases:
        f.write(f"{case['a']} {case['b']} {case['expected']}\n")

print("已生成100个混合精度测试用例到 fp16_to_fp32_mult_cases.txt")